#include "memory.h"
#include "sctlr.h"
#include "tlb.h"


static const char* IPS_NAME[7] = {
    "IPS_32BIT_4GB",
    "IPS_36BIT_64GB",
    "IPS_40BIT_1TB",
    "IPS_42BIT_4TB",
    "IPS_44BIT_16TB",
    "IPS_48BIT_256TB",
    "IPS_52BIT_4PB",
};

void config_memory_attributes(uint64_t index, uint64_t value){
    if (index > 7){
        printk("MAIR_EL! just support 8 attribute.\n");
        return;
    }
    uint64_t mair = sys_read_MAIR_EL1();
    mair |= ((value & 0xFF) << (8 * index));
    sys_write_MAIR_EL1(mair);
}


void set_page(){
    //clear page table register
    sys_write_TTBR0_EL1(0);
    sys_write_TTBR1_EL1(0);

    uint64_t mair = sys_read_MAIR_EL1();

    config_memory_attributes(MM_DEVICE_G_R_E_INDEX, MM_DEVICE_G_R_E);
    config_memory_attributes(MM_DEVICE_NG_R_E_INDEX, MM_DEVICE_NG_R_E);
    config_memory_attributes(MM_DEVICE_NG_NR_E_INDEX, MM_DEVICE_NG_NR_E);
    config_memory_attributes(MM_DEVICE_NG_NR_NE_INDEX, MM_DEVICE_NG_NR_NE);
    config_memory_attributes(MM_NORMAL_CACHEABLE_INDEX, MM_NORMAL_CACHEABLE);
    config_memory_attributes(MM_NORMAL_NOCACHEABLE_INDEX, MM_NORMAL_NOCACHEABLE);
    config_memory_attributes(6, MM_NORMAL_CACHEABLE);
    config_memory_attributes(7, MM_NORMAL_CACHEABLE);

    mair = sys_read_MAIR_EL1();

    printk("mair = %b\n", mair);

    union MemoryModelFeatureRegisterEl1 mmfr;

    mmfr.val = sys_read_ID_AA64MMFR0_EL1();

    if( mmfr.reg.TGran4 == 0xF){
        printk("4GB granukle not supported\n");
    }

    union TranslationControlRegister tcr_el1;
    tcr_el1.val = 0;
    tcr_el1.reg.IPS = mmfr.reg.PARange;

    printk("IPS: %s support \n", IPS_NAME[mmfr.reg.PARange]);

    tcr_el1.reg.Reserved1 = 0b0;
    tcr_el1.reg.DS = 0;
    
    if (mmfr.reg.ASIDBits == ASID_8_BIT){
        tcr_el1.reg.AS = 0b0;
    }
    else if (mmfr.reg.ASIDBits == ASID_16_BIT){
        tcr_el1.reg.AS = 0b1;
    }

    tcr_el1.reg.TG1 = GRANULE_SIZE_4KB;
    tcr_el1.reg.SH1 = SHAREABLITY_NON_SHAREABLE;
    tcr_el1.reg.ORGN1 = CACHEABILITY_OUTER_NON_CACHEABLE;
    tcr_el1.reg.IRGN1 = CACHEABILITY_INNER_NON_CACHEABLE;
    tcr_el1.reg.E0PD1 = 0;
    tcr_el1.reg.A1 = 0b0;
    tcr_el1.reg.T1SZ = 0b10000;

    tcr_el1.reg.TG0 = GRANULE_SIZE_4KB;
    tcr_el1.reg.SH0 = SHAREABLITY_NON_SHAREABLE;
    tcr_el1.reg.ORGN0 = CACHEABILITY_OUTER_NON_CACHEABLE;
    tcr_el1.reg.IRGN0 = CACHEABILITY_INNER_NON_CACHEABLE;
    tcr_el1.reg.E0PD0 = 0;
    tcr_el1.reg.T0SZ = 0b10000;

    sys_write_TCR_EL1(tcr_el1.val);

    // alloc page table
    uint64_t l0 = boot_mm_alloc();
    printk("l0: %x \n", l0);

    uint64_t l1 = boot_mm_alloc();
    printk("l1: %x \n", l1);

    union PageTableEntry ptbd;
    ptbd.val = 0;
	ptbd.reg.UXN = 0b0;
	ptbd.reg.PXN = 0b0;
	ptbd.reg.contiguous = 0b0;
	ptbd.reg.DBM = 0b0;
	ptbd.reg.nG = 0b1;
	ptbd.reg.AF = 0b1;
	ptbd.reg.SH = SHAREABLITY_NON_SHAREABLE;
	ptbd.reg.AP = 0b00;
	ptbd.reg.NS = 0b1;
	ptbd.reg.attrIndx = MM_DEVICE_NG_NR_NE_INDEX;
	ptbd.reg.type = PAGE_TABLE_ENTRY_TYPE_BLOCK;
	ptbd.reg.valid = PAGE_TABLE_ENTRY_VALID;

    for (uint32_t i = 0; i < 4; i++) {
		ptbd.reg.outputAddress = i; // iGB
		printk("l1-%d: %b\n", i,  ptbd.val);
		((uint64_t*)l1)[i] = ptbd.val;
	}

    union PageTableDescriptor ptd;
	ptd.val = 0;
	ptd.reg.NSTable = 0b1;
	ptd.reg.APTable = 0b1;
	ptd.reg.XNTable = 0b0;
	ptd.reg.PNXTbale = 0b0;
	ptd.reg.nextLevelTableAddr = (l1 >> PAGE_4K_OFFSET);
	ptd.reg.type = PAGE_TABLE_ENTRY_TYPE_TABLE;
	ptd.reg.valid = PAGE_TABLE_ENTRY_VALID;

	printk("l0:   0x%x\n", ptd.val);
	((uint64_t*)l0)[0] = ptd.val;

	sys_write_TTBR0_EL1(l0);
	sys_write_TTBR1_EL1(l0);
	printk("mm: identity map done.\n");

}



static void mmu_disable() {
	/*
	 * invalidate all caches
	 */
	arch_tlb_invalidate_all();

	union SystemControlRegisterEL1 sctlr;
	sctlr.val = sys_read_SCTLR_EL1();
	sctlr.reg.M = 0b0;
	sys_write_SCTLR_EL1(sctlr.val);
	dsb();
	isb();
	printk("mmu: disable\n");
}

static void mmu_enable() {
	/*
	 * invalidate all caches
	 */
	arch_tlb_invalidate_all();

	/*
	 * enable MMU
	 */
	union SystemControlRegisterEL1 sctlr;
	sctlr.val = 0;
	sctlr.reg.M = 0b1;
	sctlr.reg.C = 0b1;
	sctlr.reg.I = 0b1;
	sys_write_SCTLR_EL1(sctlr.val);
	dsb();
	isb();
	printk("mmu: enable\n");
}

void memory_init(){

    printk("------  start print momery  --------\n");
    boot_mm_init();
	mmu_disable();
    set_page();
	mmu_enable();
    printk("------  end print momery  --------\n");

}


